# pip install zhipuai 请先在终端进行安装
from zhipuai import ZhipuAI
import dashscope
import time
import os
from dotenv import load_dotenv; load_dotenv()

default_context = "你好，我是小明，我正在使用zhipuai的聊天模型进行对话。"
default_prompt = "你好，小明，今天天气怎么样？"

def zhipuai_chat_completion(context: str, prompt: str, model: str = "glm-4", temperature: float = 0.95) -> str:
	"""
	通过zhipuai调用glm模型进行聊天
	context: 聊天的上下文
	prompt: 聊天的提示语
	model: 模型名称, 默认为glm-4
	"""
	client = ZhipuAI(api_key=os.environ["ZHIPU_API_KEY"])
	error_count = 0
	start_time = time.time()
	while True:
		try:
			response = client.chat.completions.create(
				model=model,
				messages=[
					{"role": "system", "content": "你是一个财务专家，你的任务是从上下文中提取财务信息，然后回答问题。"},
					{ "role": "user", "content": context + "\n\n" + prompt },
				],
				# top_p=0.1,
				temperature=temperature,
				max_tokens=8192,
				stream=False,
				# stream=True,
			)
			break
		except Exception as e:
			print(f"AI模型调用失败: [\n{e}\n]")
			if error_count > 3:
				raise RuntimeError("已超过重试次数，AI模型调用失败")
			error_count += 1
			print(f"等待{10*error_count}秒后重试..., 重试次数: {error_count}")
			time.sleep(10*error_count)
			continue
	# print(f"AI模型调用结果: {response.json()}")
	end_time = time.time()
	# 显示耗时，格式化到秒
	print(f"耗时：{end_time - start_time}")
	print(f"本次访问{model}消耗tokens: {response.usage}")
	result = response.choices[0].message.content
	return result

#
def qwen_chat_completion(context: str = default_context, prompt: str = default_prompt, model: str = "qwen1.5-72b-chat", temperature: float = 0.85) -> str:
	"""
	通过tongyi调用glm模型进行聊天
	context: 聊天的上下文
	prompt: 聊天的提示语
	model: 模型名称, 默认为tongyi-3
	"""
	dashscope.api_key = os.environ["DASHSCOPE_API_KEY"]
	error_count = 0
	start_time = time.time()
	messages = [
		{"role": "system", "content": "你是一个财务专家，你的任务是从上下文中提取财务信息，然后回答问题。"},
		{'role': 'user', 'content': context + "\n\n" + prompt}]
	while True:
		try:
			response = dashscope.Generation.call(
				model,
				messages=messages,
				temperature=temperature,
				seed=1234,
				max_tokens=2000, # 输出最大长度
				result_format='message',  # set the result is message format.
				api_key=dashscope.api_key
			)
			if response.status_code == 200:
				break
			else:
				print('Request id: %s, Status code: %s, error code: %s, error message: %s' % (
					response.request_id, response.status_code,
					response.code, response.message
				))
				if error_count > 3:
					raise RuntimeError("已超过重试次数，AI模型调用失败")
				error_count += 1
				print(f"等待{10*error_count}秒后重试..., 重试次数: {error_count}")
				time.sleep(10*error_count)
				continue
		except Exception as e:
			print(f"AI模型调用失败: [\n{e}\n]")
			if error_count > 3:
				raise RuntimeError("已超过重试次数，AI模型调用失败")
			error_count += 1
			print(f"等待{10*error_count}秒后重试..., 重试次数: {error_count}")
			time.sleep(10*error_count)
			continue

	print(response)
	end_time = time.time()
	print(f"耗时：{end_time - start_time}")
	print(f"本次访问{model}消耗tokens: {response.usage}")
	result = response.output.choices[0].message.content
	return result

def get_context(context: str = default_context) -> str:
	"""
	获取聊天的上下文
	"""
	return context

def get_prompt(prompt: str = default_prompt) -> str:
	"""
	获取聊天的提示语
	prompt: 聊天的提示语
	"""
	return prompt

if __name__ == "__main__":
	# 需要编辑html文件将<td ...>中的属性清空，否则太长, :%s/<td[^>]*>/<td>/
	context = """
 2022年12月--部门损益汇总表【二级部】																																										
单位：元																																							操作人：	郑春	操作时间：	###############		
所属二级部	项目收入(脱税)-全年预算	项目收入(脱税)-本年累计实际	项目收入(脱税)-预算完成比	项目利润-全年预算	项目利润-本年累计实际	项目利润-预算完成比	项目毛利率-全年预算(%)	项目毛利率-本年累计实际(%)	项目毛利率-完成差值	非项目费用合计-全年预算	非项目费用合计-本年累计实际	非项目费用合计-预算完成比	售前项目-全年预算	售前项目-本年累计实际	售前项目-预算完成比	研发项目-全年预算	研发项目-本年累计实际	研发项目-预算完成比	内部管理-全年预算	内部管理-本年累计实际	内部管理-预算完成比	运营类项目-全年预算	运营类项目-本年累计实际	运营类项目-预算完成比	部门管理(人工&直接费用)-全年预算	部门管理(人工&直接费用)-本年累计实际	部门管理(人工&直接费用)-预算完成比	部门管理(R1分摊费用)-全年预算	部门管理(R1分摊费用)-本年累计实际	部门管理(R1分摊费用)-预算完成比	部门闲置-全年预算	部门闲置-本年累计实际	部门闲置-预算完成比	部门休假-全年预算	部门休假-本年累计实际	部门休假-预算完成比	部门利润(R1)-全年预算	部门利润(R1)-本年累计实际	部门利润(R1)-预算完成比	部门人数	部门人均利润(R1)	递延成本		
技术平台产品部本级	17,880,000.00	3,720,901.72	20.81%	11,000,900.00	3,720,901.72	33.82%	61.53	100.00	38.47	25,102,848.07	5,787,501.16	23.06%	0.00	318,658.74	100.00%	17,767,624.31	161,919.25	0.91%	94,000.00	46,000.00	48.94%	0.00	0.00	　	5,947,607.50	4,655,330.82	78.27%	1,116,316.26	538,339.53	48.22%	33,000.00	0.00	0.00%	144,300.00	67,252.82	46.61%	-14,101,948.07	-2,066,599.44	14.65%	3	-688,866.48	0.00		
基础开发平台产品部	0.00	4,443,758.83	100.00%	0.00	3,699,201.97	100.00%	100.00	83.24	-16.76	0.00	7,570,915.74	100.00%	0.00	0.00	　	0.00	7,069,888.93	100.00%	0.00	0.00	　	0.00	0.00	　	0.00	300,684.83	100.00%	0.00	132,684.84	100.00%	0.00	0.00	　	0.00	67,657.14	100.00%	0.00	-3,871,713.77	-100.00%	27	-143,396.81	71,108.02		
公共支撑组	0.00	0.00	　	0.00	0.00	　	100.00	100.00	0.00	0.00	2,264,468.65	100.00%	0.00	0.00	　	0.00	2,232,950.05	100.00%	0.00	0.00	　	0.00	0.00	　	0.00	136.70	100.00%	0.00	24,548.43	100.00%	0.00	0.00	　	0.00	6,833.47	100.00%	0.00	-2,264,468.65	-100.00%	10	-226,446.87	0.00		
 汇总	17,880,000.00	8,164,660.55	45.66%	11,000,900.00	7,420,103.69	67.45%	61.53	90.88	29.35	25,102,848.07	15,622,885.55	62.24%	0.00	318,658.74	100.00%	17,767,624.31	9,464,758.23	53.27%	94,000.00	46,000.00	48.94%	0.00	0.00	　	5,947,607.50	4,956,152.35	83.33%	1,116,316.26	695,572.80	62.31%	33,000.00	0.00	0.00%	144,300.00	141,743.43	98.23%	-14,101,948.07	-8,202,781.86	58.17%	40	-205,069.55	71,108.02
"""
	context_with_html = """
<table border=0 cellpadding=0 cellspacing=0 width=4971 class=xl6426884
 style='border-collapse:collapse;table-layout:fixed;width:3739pt'>
 <col class=xl6426884 width=109 span=43 style='mso-width-source:userset;
 mso-width-alt:3477;width:82pt'>
 <col class=xl6426884 width=175 style='mso-width-source:userset;mso-width-alt:
 5600;width:131pt'>
 <col class=xl6426884 width=109 style='mso-width-source:userset;mso-width-alt:
 3477;width:82pt'>
 <tr height=16 style='height:12.25pt'>
  <td>2022年12月--部门损益汇总表【二级部】</td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
 </tr>
 <tr height=16 style='height:12.25pt'>
  <td>单位：元</td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td>操作人：</td>
  <td>郑春</td>
  <td>操作时间：</td>
  <td>###############</td>
  <td></td>
  <td></td>
 </tr>
 <tr class=xl6726884 height=49 style='height:36.75pt'>
  <td>所属二级部</td>
  <td>项目收入(脱税)-全年预算</td>
  <td>项目收入(脱税)-本年累计实际</td>
  <td>项目收入(脱税)-预算完成比</td>
  <td>项目利润-全年预算</td>
  <td>项目利润-本年累计实际</td>
  <td>项目利润-预算完成比</td>
  <td>项目毛利率-全年预算(%)</td>
  <td>项目毛利率-本年累计实际(%)</td>
  <td>项目毛利率-完成差值</td>
  <td>非项目费用合计-全年预算</td>
  <td>非项目费用合计-本年累计实际</td>
  <td>非项目费用合计-预算完成比</td>
  <td>售前项目-全年预算</td>
  <td>售前项目-本年累计实际</td>
  <td>售前项目-预算完成比</td>
  <td>研发项目-全年预算</td>
  <td>研发项目-本年累计实际</td>
  <td>研发项目-预算完成比</td>
  <td>内部管理-全年预算</td>
  <td>内部管理-本年累计实际</td>
  <td>内部管理-预算完成比</td>
  <td>运营类项目-全年预算</td>
  <td>运营类项目-本年累计实际</td>
  <td>运营类项目-预算完成比</td>
  <td>部门管理(人工&amp;直接费用)-全年预算</td>
  <td>部门管理(人工&amp;直接费用)-本年累计实际</td>
  <td>部门管理(人工&amp;直接费用)-预算完成比</td>
  <td>部门管理(R1分摊费用)-全年预算</td>
  <td>部门管理(R1分摊费用)-本年累计实际</td>
  <td>部门管理(R1分摊费用)-预算完成比</td>
  <td>部门闲置-全年预算</td>
  <td>部门闲置-本年累计实际</td>
  <td>部门闲置-预算完成比</td>
  <td>部门休假-全年预算</td>
  <td>部门休假-本年累计实际</td>
  <td>部门休假-预算完成比</td>
  <td>部门利润(R1)-全年预算</td>
  <td>部门利润(R1)-本年累计实际</td>
  <td>部门利润(R1)-预算完成比</td>
  <td>部门人数</td>
  <td>部门人均利润(R1)</td>
  <td>递延成本</td>
  <td></td>
  <td></td>
 </tr>
 <tr class=xl6726884 height=33 style='height:24.5pt'>
  <td>技术平台产品部本级</td>
  <td>17,880,000.00</td>
  <td>3,720,901.72</td>
  <td>20.81%</td>
  <td>11,000,900.00</td>
  <td>3,720,901.72</td>
  <td>33.82%</td>
  <td>61.53</td>
  <td>100.00</td>
  <td>38.47</td>
  <td>25,102,848.07</td>
  <td>5,787,501.16</td>
  <td>23.06%</td>
  <td>0.00</td>
  <td>318,658.74</td>
  <td>100.00%</td>
  <td>17,767,624.31</td>
  <td>161,919.25</td>
  <td>0.91%</td>
  <td>94,000.00</td>
  <td>46,000.00</td>
  <td>48.94%</td>
  <td>0.00</td>
  <td>0.00</td>
  <td>　</td>
  <td>5,947,607.50</td>
  <td>4,655,330.82</td>
  <td>78.27%</td>
  <td>1,116,316.26</td>
  <td>538,339.53</td>
  <td>48.22%</td>
  <td>33,000.00</td>
  <td>0.00</td>
  <td>0.00%</td>
  <td>144,300.00</td>
  <td>67,252.82</td>
  <td>46.61%</td>
  <td>-14,101,948.07</td>
  <td>-2,066,599.44</td>
  <td>14.65%</td>
  <td>3</td>
  <td>-688,866.48</td>
  <td>0.00</td>
  <td></td>
  <td></td>
 </tr>
 <tr class=xl6726884 height=33 style='height:24.5pt'>
  <td>基础开发平台产品部</td>
  <td>0.00</td>
  <td>4,443,758.83</td>
  <td>100.00%</td>
  <td>0.00</td>
  <td>3,699,201.97</td>
  <td>100.00%</td>
  <td>100.00</td>
  <td>83.24</td>
  <td>-16.76</td>
  <td>0.00</td>
  <td>7,570,915.74</td>
  <td>100.00%</td>
  <td>0.00</td>
  <td>0.00</td>
  <td>　</td>
  <td>0.00</td>
  <td>7,069,888.93</td>
  <td>100.00%</td>
  <td>0.00</td>
  <td>0.00</td>
  <td>　</td>
  <td>0.00</td>
  <td>0.00</td>
  <td>　</td>
  <td>0.00</td>
  <td>300,684.83</td>
  <td>100.00%</td>
  <td>0.00</td>
  <td>132,684.84</td>
  <td>100.00%</td>
  <td>0.00</td>
  <td>0.00</td>
  <td>　</td>
  <td>0.00</td>
  <td>67,657.14</td>
  <td>100.00%</td>
  <td>0.00</td>
  <td>-3,871,713.77</td>
  <td>-100.00%</td>
  <td>27</td>
  <td>-143,396.81</td>
  <td>71,108.02</td>
  <td></td>
  <td></td>
 </tr>
 <tr class=xl6726884 height=16 style='height:12.25pt'>
  <td>公共支撑组</td>
  <td>0.00</td>
  <td>0.00</td>
  <td>　</td>
  <td>0.00</td>
  <td>0.00</td>
  <td>　</td>
  <td>100.00</td>
  <td>100.00</td>
  <td>0.00</td>
  <td>0.00</td>
  <td>2,264,468.65</td>
  <td>100.00%</td>
  <td>0.00</td>
  <td>0.00</td>
  <td>　</td>
  <td>0.00</td>
  <td>2,232,950.05</td>
  <td>100.00%</td>
  <td>0.00</td>
  <td>0.00</td>
  <td>　</td>
  <td>0.00</td>
  <td>0.00</td>
  <td>　</td>
  <td>0.00</td>
  <td>136.70</td>
  <td>100.00%</td>
  <td>0.00</td>
  <td>24,548.43</td>
  <td>100.00%</td>
  <td>0.00</td>
  <td>0.00</td>
  <td>　</td>
  <td>0.00</td>
  <td>6,833.47</td>
  <td>100.00%</td>
  <td>0.00</td>
  <td>-2,264,468.65</td>
  <td>-100.00%</td>
  <td>10</td>
  <td>-226,446.87</td>
  <td>0.00</td>
  <td></td>
  <td></td>
 </tr>
 <tr class=xl6726884 height=16 style='height:12.25pt'>
  <td><span style='mso-spacerun:yes'> </span>汇总</td>
  <td>17,880,000.00</td>
  <td>8,164,660.55</td>
  <td>45.66%</td>
  <td>11,000,900.00</td>
  <td>7,420,103.69</td>
  <td>67.45%</td>
  <td>61.53</td>
  <td>90.88</td>
  <td>29.35</td>
  <td>25,102,848.07</td>
  <td>15,622,885.55</td>
  <td>62.24%</td>
  <td>0.00</td>
  <td>318,658.74</td>
  <td>100.00%</td>
  <td>17,767,624.31</td>
  <td>9,464,758.23</td>
  <td>53.27%</td>
  <td>94,000.00</td>
  <td>46,000.00</td>
  <td>48.94%</td>
  <td>0.00</td>
  <td>0.00</td>
  <td>　</td>
  <td>5,947,607.50</td>
  <td>4,956,152.35</td>
  <td>83.33%</td>
  <td>1,116,316.26</td>
  <td>695,572.80</td>
  <td>62.31%</td>
  <td>33,000.00</td>
  <td>0.00</td>
  <td>0.00%</td>
  <td>144,300.00</td>
  <td>141,743.43</td>
  <td>98.23%</td>
  <td>-14,101,948.07</td>
  <td>-8,202,781.86</td>
  <td>58.17%</td>
  <td>40</td>
  <td>-205,069.55</td>
  <td>71,108.02</td>
  <td></td>
  <td></td>
 </tr>
 <![if supportMisalignedColumns]>
 <tr height=0 style='display:none'>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
  <td></td>
 </tr>
 <![endif]>
</table>
"""
	# prompt = "根据上述内容，提取所有部门一级标题为\"部门管理(R1分摊费用)\"项下的所有栏位数据，输出采用Json格式"
	# prompt = "根据上述内容，提取各部门\"部门管理(R1分摊费用)\"下所有栏位数据，输出采用Json格式"
	prompt = '根据上述内容，提取所有部门"部门管理(R1分摊费用)-本年累计实际"下的数据，输出采用Json格式，{"部门名称": {"key1": "value1"},}'

	result = zhipuai_chat_completion(context, prompt, model="glm-4", temperature = 0.1)

	# result = qwen_chat_completion(context, prompt, model="qwen1.5-72b-chat", temperature = 0.85)

	# result = qwen_chat_completion(context, prompt, model="qwen-plus")
	print(f"prompt:{prompt}")
	print(f"result:{result}")
